import torch
import numpy as np
import os
import random

class BaseSampleGeneration:
    """
    生成的方式，尽量多的探索一下，接下来一段时间的重点，是找基样本的生成方法，目前想到的方法有浮点数法，方向法。都试一试，再试试少量像素点的攻击方法。
    """

    def __init__(self, datater, neter, config):

        self.dataer = datater
        self.network = neter
        self.config = config
        self.path = './data/BaseSample'
        self.path = os.path.join(self.path, self.dataer.dataset_name)        
        if os.path.exists(self.path) == False:
            os.mkdir(self.path)


    def get_BaseSample(self, class_index):

        bs_path = os.path.join(self.path, 'BS_sample_size{}_method{}'.format(self.config['bs_size'], self.config['bs_method']))
        if os.path.exists(bs_path) == False:
            self.generate_BaseSample()
        print('Loading Base Sample from <{}>'.format(bs_path))
        bs_samples = torch.load(bs_path)
        print('Load Done.')
        return bs_samples[class_index]

    def generate_BaseSample(self, class_index_list=[], isAll=True):

        if isAll:
            class_index_list = [i for i in range(self.dataer.class_num)]

        bs_samples = {}

        print('=='*50)
        print('Generating Base Samples, details as follows:')
        print('Dataset: {}'.format(self.dataer.dataset_name))
        print('Base Samples Size: {}'.format(self.config['bs_size']))   
        print('Base Samples Method: {}'.format(self.config['bs_method']))
        print('Total Base Sample Class Number: {}'.format(len(class_index_list)))   
        print('=='*50)
        print()
        print('hehe')
        for class_index in class_index_list:
            print('Generating Class {} Base Samples for Dataset {}'.format(class_index, self.dataer.dataset_name), end='  ')
            print('huhu')
            loader = self.dataer.get_class_loader(class_index=class_index)            
            count = 0
            bs_samples[class_index] = []
            print('count')
            for sample, label in loader:
                base_sample, flag = self.evolutionary_search(sample, label)
                if flag:
                    count += 1
                    bs_samples[class_index].append(base_sample)
                if count == self.config['bs_size']:
                    break
            if count == 0:
                print('undone.')
                raise Exception('Can not find the base sample for class index {}'.format(class_index))
            elif count < self.config['bs_size']:
                print('find {} base samples'.format(len(bs_samples[class_index])))
            else:
                print('done.')
            bs_samples[class_index] = torch.cat(bs_samples[class_index], dim=0)

        bs_path = os.path.join(self.path, 'BS_sample_size{}_method{}'.format(self.config['bs_size'], self.config['bs_method']))
        print('Saving Base Sample to <{}>'.format(bs_path))
        torch.save(bs_samples, bs_path)
        print('Save Done.')

    def evolutionary_search(self, sample, label, norm_type='inf', cof=0.2):
        """
        Execute evolutionary search strategy to obtain base samples
        return: low norm base sample(Tensor), isSuccessful
        """
        def get_fitness(samples, label):

            probability_output = self.network.get_prediction_score(samples)
            fitness = -(probability_output[:, label].squeeze()).detach().cpu()
            fitness += 3 * torch.linalg.vector_norm(samples, ord=float('inf'), dim=(1, 2, 3))
            return fitness
        
        def judge(samples, label, proba_limit=0.99, norm_limit=0.05):

            probability_output = self.network.get_prediction_score(samples)
            norms = torch.linalg.vector_norm(samples, ord=float('inf'), dim=(1, 2, 3))
            
            for i in range(samples.shape[0]):
                print('population: {}, probability: {}, norm: {}'.format(i, probability_output[i, label], norms[i]))
                if probability_output[i, label] >= proba_limit and norms[i] < norm_limit:
                    return i
            return -1

        population = sample + (-cof + 2 * cof * torch.rand(self.config['bs_pop_size'], sample.shape[1], sample.shape[2], sample.shape[3]))
        population = torch.clamp(population, min=0., max=1.)
        fitness = get_fitness(population, label)

        generations = 0
        while generations < 100:

            ## mutation
            Mpopulation = torch.zeros_like(population)
            for i in range(population.shape[0]):
                r1 = r2 = r3 = 0
                while r1 == i or r2 == i or r3 == i or r3 == r1 or r3 == r2:
                    r1 = random.randint(0, population.shape[0] - 1)
                    r2 = random.randint(0, population.shape[0] - 1)
                    r3 = random.randint(0, population.shape[0] - 1)
                Mpopulation[i] = population[r1] + self.config['F'] * (population[r2] - population[r3])
            Mpopulation = torch.clamp(Mpopulation, min=0., max=1.)

            ## cross
            Cpopulation = torch.zeros_like(population)
            for i in range(population.shape[0]):
                dim = np.prod(population.shape[1:])
                for j in range(dim):
                    channel = j // np.prod(population.shape[2:])
                    x = (j % np.prod(population.shape[2:])) // np.prod(population.shape[3:])
                    y = (j % np.prod(population.shape[2:])) % np.prod(population.shape[3:])
                    rand_j = random.randint(0, dim - 1)
                    rand_float = random.random()
                    if rand_float < self.config['CR'] or rand_j == j:
                        Cpopulation[i, channel, x, y] = Mpopulation[i, channel, x, y]
                    else:
                        Cpopulation[i, channel, x, y] = population[i, channel, x, y]

            ## selection
            cfitness = get_fitness(Cpopulation, label)
            for i in range(population.shape[0]):
                if cfitness[i] < fitness[i]:
                    population[i] = Cpopulation[i]
                    fitness[i] = cfitness[i]

            bs_index = judge(population, label)
            if bs_index != -1:
                return population[bs_index], True

            generations += 1
            print('generation : {}'.format(generations))

        return None, False